#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Time    : 2025/4/1 15:27
# @Author  : 梁满仓
# @File    : breakout_strategy_ma7.py
# @Software: PyCharm

import akshare as ak
import pandas as pd
import numpy as np
from datetime import datetime


# 获取股票数据
def get_stock_data(stock_code, start_date, end_date):
    """
    获取指定股票的历史行情数据
    :param stock_code: 股票代码，如"000001"
    :param start_date: 开始日期，格式"YYYYMMDD"
    :param end_date: 结束日期，格式"YYYYMMDD"
    :return: DataFrame
    """
    #start = datetime.strptime(start_date, "%Y%m%d").strftime("%Y-%m-%d")
    #end = datetime.strptime(end_date, "%Y%m%d").strftime("%Y-%m-%d")

    #df = ak.stock_zh_a_hist(symbol=stock_code, period="daily", start_date=start, end_date=end, adjust="hfq")
    df = ak.stock_zh_a_hist(symbol=stock_code, period="daily", start_date=start_date, end_date=end_date, adjust="hfq")

    df['日期'] = pd.to_datetime(df['日期'])
    df = df.set_index('日期').sort_index()
    return df


# 计算交易信号（修改为当日收盘价交易）
def calculate_signals(df, ma_period=7):
    """
    计算买入和卖出信号（当日收盘价版本）
    :param df: 包含股票数据的DataFrame
    :param ma_period: 均线周期，默认为7
    :return: 添加了信号的DataFrame
    """
    # 计算7日均线
    df['MA7'] = df['收盘'].rolling(window=ma_period).mean()

    # 前一日收盘价和均线值
    df['prev_close'] = df['收盘'].shift(1)
    df['prev_MA7'] = df['MA7'].shift(1)

    # 突破信号（当日收盘价上穿7日均线）
    df['buy_signal'] = (df['收盘'] > df['MA7']) & (df['prev_close'] <= df['prev_MA7'])

    # 跌破信号（当日收盘价下穿7日均线）
    df['sell_signal'] = (df['收盘'] < df['MA7']) & (df['prev_close'] >= df['prev_MA7'])

    return df


# 执行回测（修改为当日收盘价交易）
def backtest(df):
    """
    执行回测并生成交易记录（当日收盘价版本）
    :param df: 包含信号的DataFrame
    :return: 交易记录DataFrame和统计信息
    """
    trades = []
    position = None  # 当前持仓

    for date, row in df.iterrows():
        if position is None and row['buy_signal']:
            # 买入信号 - 以当日收盘价买入
            position = {
                'buy_date': date,
                'buy_price': row['收盘'],
                'sell_date': None,
                'sell_price': None
            }
        elif position is not None and row['sell_signal']:
            # 卖出信号 - 以当日收盘价卖出
            position['sell_date'] = date
            position['sell_price'] = row['收盘']
            trades.append(position)
            position = None

    # 如果最后还有持仓，以最后一天收盘价平仓
    if position is not None:
        position['sell_date'] = df.index[-1]
        position['sell_price'] = df.iloc[-1]['收盘']
        trades.append(position)

    trades_df = pd.DataFrame(trades)

    # 计算每笔交易的收益率和持有天数
    if not trades_df.empty:
        trades_df['return'] = (trades_df['sell_price'] - trades_df['buy_price']) / trades_df['buy_price']
        trades_df['holding_days'] = (trades_df['sell_date'] - trades_df['buy_date']).dt.days

    return trades_df


# 计算统计信息
def calculate_stats(trades_df):
    """
    计算回测统计信息
    :param trades_df: 交易记录DataFrame
    :return: 统计信息字典
    """
    if trades_df.empty:
        return {
            "总交易次数": 0,
            "盈利交易比例": 0,
            "平均收益率": 0,
            "最大收益率": 0,
            "最小收益率": 0,
            "总收益率": 0,
            "平均持有天数": 0
        }

    stats = {
        "总交易次数": len(trades_df),
        "盈利交易比例": len(trades_df[trades_df['return'] > 0]) / len(trades_df),
        "平均收益率": trades_df['return'].mean(),
        "最大收益率": trades_df['return'].max(),
        "最小收益率": trades_df['return'].min(),
        "总收益率": trades_df['return'].sum(),
        "平均持有天数": trades_df['holding_days'].mean()
    }

    return stats


# 主函数
def main(stock_code, start_date, end_date):
    # 获取数据
    df = get_stock_data(stock_code, start_date, end_date)

    # 计算信号
    df = calculate_signals(df)

    # 执行回测
    trades_df = backtest(df)

    # 计算统计信息
    stats = calculate_stats(trades_df)

    return df, trades_df, stats


# 示例使用
if __name__ == "__main__":
    stock_code = "002128"  # 股票代码
    start_date = "20220101"
    end_date = "20221231"

    df, trades_df, stats = main(stock_code, start_date, end_date)

    print(f"\n{stock_code} 交易记录:")
    print(trades_df)
    trades_df.to_excel(f"{stock_code}_7日突破均线记录.xlsx", index=False)

    print("\n收益统计:")
    for key, value in stats.items():
        if isinstance(value, float):
            print(f"{key}: {value:.2%}")
        else:
            print(f"{key}: {value}")
